import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.express as px
from io import BytesIO

# 设置应用的标题和描述
st.title('中国各省新冠病毒感染数据分析平台')
st.write('请上传包含各城市新冠病毒感染数据的CSV文件，并进行数据分析和可视化。')

# 上传CSV文件
uploaded_file = st.file_uploader("上传CSV文件", type=["csv"])

# 初始化数据框
df = pd.DataFrame()

# 当文件被上传时，尝试读取数据
if uploaded_file is not None:
    try:
        df = pd.read_csv(uploaded_file)

        # 检查必要的列是否存在
        required_columns = [
            'id', 'confirmedCount', 'confirmedIncr', 'curedCount', 'curedIncr',
            'currentConfirmedCount', 'currentConfirmedIncr', 'dateId', 'deadCount',
            'deadIncr', 'suspectedCount', 'suspectedCountIncr'
        ]
        if not all(column in df.columns for column in required_columns):
            st.error("CSV文件缺少必要的列，请检查文件格式。")
        else:
            # 数据预处理
            # 去除空值
            df.dropna(inplace=True)

            # 异常值处理
            df.replace({k: -1 for k in df.columns if df[k].dtype == 'int64'}, np.nan, inplace=True)
            df.dropna(inplace=True)

            # 将日期列转换为日期类型，并按日期排序
            df['dateId'] = pd.to_datetime(df['dateId'])
            df.sort_values('dateId', inplace=True)

            # 用户交互：时间范围选择
            start_date = pd.to_datetime(st.date_input('选择开始日期', value=df['dateId'].min().to_pydatetime().date()))
            end_date = pd.to_datetime(st.date_input('选择结束日期', value=df['dateId'].max().to_pydatetime().date()))

            # 创建布尔掩码来过滤日期范围内的数据
            mask = (df['dateId'] >= start_date) & (df['dateId'] <= end_date)
            filtered_df = df.loc[mask]

            # 用户交互：阈值输入
            threshold = st.number_input('设置感染人数阈值', min_value=0, value=100)

            # 数据分析：计算总的新冠病毒感染趋势
            total_infected = filtered_df.groupby('dateId')['confirmedCount'].sum().reset_index()

            # 图表展示：总感染趋势图
            st.subheader('总感染趋势图')
            fig, ax = plt.subplots(figsize=(10, 6))  # 调整图表大小
            sns.lineplot(x='dateId', y='confirmedCount', data=total_infected, ax=ax)
            ax.set_title('总感染趋势(sum)')
            ax.set_xlabel('日期')
            ax.set_ylabel('累计确诊人数')
            st.pyplot(fig)

            # 数据分析：对每个城市，计算感染人数随时间变化的曲线图
            city_fig = px.line(filtered_df, x='dateId', y='confirmedIncr', color='id', title='各城市感染曲线图')
            st.plotly_chart(city_fig, use_container_width=True)  # 使图表适应容器宽度

            # 风险区域识别：找出高风险城市
            high_risk_cities = filtered_df[filtered_df['confirmedCount'] > threshold]['id'].unique()

            # 图表展示：高风险区域图
            if len(high_risk_cities) > 0:
                st.subheader('高风险区域感染趋势')
                high_risk_df = filtered_df[filtered_df['id'].isin(high_risk_cities)]
                high_risk_fig = px.line(high_risk_df, x='dateId', y='confirmedCount', color='id',title='高风险区域感染趋势')
                st.plotly_chart(high_risk_fig, use_container_width=True)
            else:
                st.write('没有城市超过设定的感染人数阈值。')

            # 结果输出：下载图表功能
            def save_plot_to_bytes(figure):
                buf = BytesIO()
                figure.savefig(buf, format='png')
                buf.seek(0)
                return buf

            # 下载总感染趋势图
            st.subheader('下载图表')
            download_link = st.button('下载总感染趋势图')
            if download_link:
                fig, ax = plt.subplots()
                sns.lineplot(x='dateId', y='confirmedCount', data=total_infected, ax=ax)
                ax.set_title('总感染趋势')
                ax.set_xlabel('日期')
                ax.set_ylabel('累计确诊人数')
                img = save_plot_to_bytes(fig)
                st.download_button('下载图表', img, file_name='total_infected_trend.png', mime='image/png')

    except Exception as e:
        st.error(f"读取或处理文件时发生错误：{str(e)}")

# 如果没有上传文件，则显示提示信息
if df.empty:
    st.info("请上传一个包含指定字段的CSV文件以进行分析。")
